#pragma once

#include <cassert>
#include <tuple>
#include <type_traits>

#include "result.h"

namespace limbo {

namespace details {

template <typename T, typename Context>
constexpr auto has_prefix = []() -> bool {
  if constexpr (std::is_same_v<
                    decltype(T::prefix(
                        std::declval<std::add_lvalue_reference_t<Context>>())),
                    uint32_t>) {
    return true;
  } else {
    return false;
  }
};

template <typename Layer, typename TargetState, size_t I>
constexpr size_t get_index_impl() {
  static_assert(I < std::tuple_size_v<Layer>, "no state on the layer");
  using State = std::tuple_element_t<I, Layer>;
  if constexpr (std::is_same_v<TargetState, State>) {
    return I;
  } else {
    return get_index_impl<Layer, TargetState, I + 1>();
  }
}

template <typename Layer, typename State> constexpr size_t get_index() {
  return get_index_impl<Layer, State, 0>();
}

using Location = std::pair<size_t, size_t>;

template <typename Layers, typename TargetContext, size_t I, size_t J>
constexpr Location locate_context() {
  using Layer = std::tuple_element_t<I, Layers>;
  using State = std::tuple_element_t<J, Layer>;
  using Context = typename State::Context;
  if constexpr (std::is_same_v<TargetContext, Context>) {
    return {I, J};
  } else {
    constexpr auto states_sz = std::tuple_size_v<Layer>;
    if constexpr (J + 1 < states_sz) {
      return locate_context<Layers, TargetContext, I, J + 1>();
    } else {
      constexpr auto layers_sz = std::tuple_size_v<Layers>;
      static_assert(I < layers_sz, "no state found in the layers");
      return locate_context<Layers, TargetContext, I + 1, 0>();
    }
  }
}

template <uint8_t sz> using sz_t = std::integral_constant<uint8_t, sz>;

template <typename Packet, typename SZ = sz_t<0>> struct trailer_size {
  static constexpr uint32_t value = SZ::value;
};
template <typename Packet>
struct trailer_size<Packet, sz_t<Packet::TRAILER_SZ>> {
  static constexpr uint32_t value = Packet::TRAILER_SZ;
};

template <typename Context> constexpr uint32_t get_trailer_size() {
  using Packet = typename Context::Packet;
  return trailer_size<Packet>::value;
}

template <typename, typename = void> constexpr bool is_packetable{};

template <typename T>
constexpr bool is_packetable<T, std::void_t<typename T::Packet>> = true;

template <typename T> static constexpr auto get_prefixable_context(T &ctx) {
  using State = typename T::ContextState;
  auto lower_ctx = get_lower_context(ctx);
  if constexpr (!is_packetable<State>) {
    return get_prefixable_context(*lower_ctx);
  } else {
    return lower_ctx;
  }
}

template <typename, typename = void> constexpr bool is_frameable{true};

template <typename T>
constexpr bool is_frameable<T, std::void_t<typename T::non_frameable>> =
    T::non_frameable::value;

} // namespace details

template <typename... Ts> using Layer = std::tuple<Ts...>;

template <typename... Ls> struct Stack {
  using Layers = std::tuple<Ls...>;
  static constexpr auto N = std::tuple_size_v<Layers>;
  using LowestLayer = std::tuple_element_t<N - 1, Layers>;
  using LowestState = std::tuple_element_t<0, LowestLayer>;
  using EntryState = typename LowestState::LowerState;
  static_assert(std::tuple_size_v<LowestLayer> == 1,
                "lowest layer have to be only one");

  static_assert(std::is_default_constructible_v<Layers>,
                "every state should be default constructible");

  template <size_t I, size_t J> auto &get() {
    auto &layer = std::get<I>(layers);
    auto &state = std::get<J>(layer);
    return state;
  }

  struct Sender {
    Sender(uint32_t error_code, void *state)
        : error_code_{error_code}, state_{state} {}

    Sender(Chunk envelope, uint32_t trailer_size)
        : error_code_{0}, envelope_{envelope}, trailer_size_{trailer_size} {
      assert(trailer_size_ < envelope_.size());
    }

    Chunk header() {
      auto sz = envelope_.size() - trailer_size_;
      return Chunk(envelope_.data(), sz);
    }

    Chunk trailer() {
      auto ptr = envelope_.data() + envelope_.size() - trailer_size_;
      return Chunk(ptr, trailer_size_);
    }

    uint32_t error_code() { return error_code_; }
    void *error_state() { return state_; }
    operator bool() const { return error_code_ == 0; }

  private:
    uint32_t error_code_;
    void *state_;
    Chunk envelope_;
    uint32_t trailer_size_;
  };

  template <typename Context>
  Result send(Chunk buff, Context &ctx, Chunk payload) {
    constexpr auto location = details::locate_context<Layers, Context, 0, 0>();
    static constexpr auto I = location.first;
    static constexpr auto J = location.second;
    auto prefixable_ctx = details::get_prefixable_context(ctx);
    auto prefix = prefix_impl(*prefixable_ctx);
    auto sub_buff = buff.skip(prefix);
    return send_impl<I, J, true, Context>(sub_buff, ctx, payload);
  }

  template <typename Context>
  Sender send_envelope(Chunk buff, Context &ctx, Chunk payload) {
    constexpr auto location = details::locate_context<Layers, Context, 0, 0>();
    static constexpr auto I = location.first;
    static constexpr auto J = location.second;
    auto prefixable_ctx = details::get_prefixable_context(ctx);
    auto prefix = prefix_impl(*prefixable_ctx);
    auto sub_buff = buff.skip(prefix);
    auto result = send_impl<I, J, false, Context>(sub_buff, ctx, payload);
    if (result) {
      auto trailer_sz = trailer_size<Context>();
      return Sender(result.consumed(), trailer_sz);
    } else {
      return Sender(result.error_code(), result.state());
    }
  }

  Result recv(Chunk buff, EntryState *entry_state) {
    auto &layer = std::get<N - 1>(layers);
    auto &state = std::get<0>(layer);
    auto result = state.recv(buff, entry_state);
    if (!result) {
      return result;
    }
    if (result.demand()) {
      return result;
    }
    if constexpr (N > 1) {
      auto inner_buff = state.get_parsed().payload;
      return recv_upper<N - 2, 0>(inner_buff, &state, result);
    }
    return result;
  }

private:
  template <size_t I, size_t J, bool CopyContent, typename Context>
  Result send_impl(Chunk buff, Context &ctx, Chunk payload) {
    using Layer = std::tuple_element_t<I, Layers>;
    using State = std::tuple_element_t<J, Layer>;
    if constexpr (details::is_packetable<State>) {
      using Packet = typename State::Packet;
      auto &layer = std::get<I>(layers);
      auto &state = std::get<J>(layer);
      auto result =
          Packet::template send<CopyContent>(state, ctx, buff, payload);
      if (!result) {
        return result;
      }
      if (result.demand()) {
        return result;
      }
      if constexpr (I + 1 < N) {
        using LowerLayer = std::tuple_element_t<I + 1, Layers>;
        using LowerState = typename State::LowerState;
        static constexpr auto K = details::get_index<LowerLayer, LowerState>();
        auto packet = result.consumed();
        auto new_payload = [&]() -> Chunk {
          if constexpr (!CopyContent) {
            packet = Chunk(packet.data(), packet.size() + payload.size());
          }
          return packet;
        }();

        auto lower_context = get_lower_context(ctx);
        auto prefix = get_prefix(*lower_context);
        auto sub_buff = buff.skip(-prefix);
        result = send_impl<I + 1, K, CopyContent>(sub_buff, *lower_context,
                                                  new_payload);
        if constexpr (!CopyContent) {
          if (result && !result.demand()) {
            auto consumed = result.consumed();
            auto sz = packet.size() + consumed.size() - payload.size();
            auto header = Chunk(consumed.data(), sz);
            return Result::make_consumed(result.state(), header);
          }
        }
      }
      return result;
    } else {
      using LowerLayer = std::tuple_element_t<I + 1, Layers>;
      using LowerState = typename State::LowerState;
      static constexpr auto K = details::get_index<LowerLayer, LowerState>();
      auto lower_ctx = get_lower_context(ctx);
      return send_impl<I + 1, K, CopyContent>(buff, *lower_ctx, payload);
    }
  }

  template <typename Context> uint32_t prefix_impl(const Context &ctx) {
    using State = typename Context::ContextState;
    auto value = get_prefix(ctx);
    if constexpr (details::has_lower_state_v<State>) {
      return value + prefix_impl(*get_lower_context(ctx));
    } else {
      return value;
    }
  }

  template <typename Context> uint32_t trailer_size() {
    using State = typename Context::ContextState;
    auto value = details::get_trailer_size<State>();
    if constexpr (details::has_lower_state_v<State>) {
      using LowerState = typename State::LowerState;
      using LowerContext = typename LowerState::Context;
      return value + trailer_size<LowerContext>();
    } else {
      return value;
    }
  }

  template <size_t I, size_t J, typename LowerState>
  Result recv_upper(Chunk buff, LowerState *lower_state, Result &lower_result) {
    using Layer = std::tuple_element_t<I, Layers>;
    using State = std::tuple_element_t<J, Layer>;
    if constexpr (!details::is_frameable<LowerState>) {
      if (!buff.size()) {
        return lower_result;
      }
    }
    if constexpr (std::is_same_v<typename State::LowerState, LowerState>) {
      auto &layer = std::get<I>(layers);
      auto &state = std::get<J>(layer);
      auto &parsed = lower_state->get_parsed();
      if (state.select(&parsed)) {
        auto result = state.recv(buff, lower_state);
        result.merge(lower_result.consumed());
        if constexpr (I) {
          if (!result || result.demand()) {
            return result;
          }
          auto inner_buff = get_payload(state.get_parsed());
          return recv_upper<I - 1, 0>(inner_buff, &state, result);
        }
        return result;
      }
    }
    if constexpr (J + 1 < std::tuple_size_v<Layer>) {
      return recv_upper<I, J + 1>(buff, lower_state, lower_result);
    }
    return lower_result;
  }

  Layers layers;
};

} // namespace limbo
